package com.api.apigateway.risk;


import ai.onnxruntime.*;
import lombok.extern.slf4j.Slf4j;

import java.io.*;
import java.nio.FloatBuffer;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;

/**
 * ONNX 推理（IsolationForest / OneClassSVM）— JDK8 兼容版
 * 运行逻辑：
 * 1) 构造函数加载 .onnx（支持文件路径或 classpath: 前缀）
 * 2) warmup() 预热一次，避免首个请求抖动
 * 3) score(fv) 构造 [1, N] 输入 → session.run → 解析第一个或指定输出 → 归一化到 [0,1]
 *
 * 注意：
 * - ORDER 必须与你训练脚本一致；
 * - 若模型输出为 decision_function（OCSVM 常见），可按需翻转/压缩到 [0,1]。
 */
@Slf4j
public class OnnxIsolationScorer implements ModelScorer, AutoCloseable {

    private final OrtEnvironment env;     // ONNX 运行时环境
    private final OrtSession session;     // 模型会话（线程安全，可复用）
    private final String[] ORDER;         // 特征顺序
    private final String outputName;      // 输出名；null 则取第一个输出
    private final boolean flipDecision;   // 是否需要对决策函数做翻转（OCSVM）

    public OnnxIsolationScorer(String modelPathOrClasspath,
                               String[] featureOrder,
                               String explicitOutputName) {
        this(modelPathOrClasspath, featureOrder, explicitOutputName, false);
    }

    public OnnxIsolationScorer(String modelPathOrClasspath,
                               String[] featureOrder,
                               String explicitOutputName,
                               boolean flipDecisionFunction) {
        try {
            this.ORDER = featureOrder.clone();
            this.outputName = explicitOutputName;
            this.flipDecision = flipDecisionFunction;

            this.env = OrtEnvironment.getEnvironment();
            byte[] modelBytes = loadBytes(modelPathOrClasspath);

            OrtSession.SessionOptions opts = new OrtSession.SessionOptions();
            // 可选：提高兼容性/性能的图优化等级
            // opts.setOptimizationLevel(OrtSession.SessionOptions.OptLevel.BASIC_OPT);

            this.session = env.createSession(modelBytes, opts);
            warmup();
        } catch (Exception e) {
            throw new RuntimeException("Init OnnxIsolationScorer failed: " + e.getMessage(), e);
        }
    }

    /** JDK8 版：文件或 classpath 资源读成 byte[] */
    private static byte[] loadBytes(String pathOrCp) throws Exception {
        InputStream in;
        if (pathOrCp.startsWith("classpath:")) {
            String cp = pathOrCp.substring("classpath:".length());
            if (!cp.startsWith("/")) cp = "/" + cp;
            in = OnnxIsolationScorer.class.getResourceAsStream(cp);
            if (in == null) throw new IllegalArgumentException("Classpath resource not found: " + cp);
        } else {
            in = new FileInputStream(new File(pathOrCp));
        }
        ByteArrayOutputStream bos = new ByteArrayOutputStream(8192);
        try {
            byte[] buf = new byte[8192];
            int n;
            while ((n = in.read(buf)) != -1) {
                bos.write(buf, 0, n);
            }
            return bos.toByteArray();
        } finally {
            try { in.close(); } catch (IOException ignore) {}
            try { bos.close(); } catch (IOException ignore) {}
        }
    }

    /** 预热：用 0 向量跑一次，避免首个请求延迟抖动 */
    private void warmup() throws Exception {
        float[] zeros = new float[ORDER.length];
        OnnxTensor t = OnnxTensor.createTensor(env, FloatBuffer.wrap(zeros), new long[]{1, ORDER.length});
        try {
            Map<String, OnnxTensorLike> input = new HashMap<String, OnnxTensorLike>(1);
            String inName = session.getInputNames().iterator().next();
            input.put(inName, t);
            OrtSession.Result r = null;
            try {
                r = session.run(input);
            } finally {
                if (r != null) r.close();
            }
        } finally {
            t.close();
        }
    }

    /**
     * Isolation Forest 模型（iforest.onnx）
     *  Isolation Forest 是一个无监督的机器学习模型，常用于检测    异常值。它的工作原理是：
     *      模型的输出是“正常度”或“异常度”。
     *      scores 输出的是一个值，这个值越高表示越正常（和模型的异常检测方向是相反的）。具体来说：
     *      scores 越接近 1 → 该请求被模型判定为“正常”。
     *      scores 越接近 0 → 该请求被模型判定为“异常”
     * @param fv 特征向量
     * @return
     */
    @Override
    public double score(FeatureVector fv) {
        try {
            // 1) 特征向量 → float[]
            double[] arr = fv.toArray(ORDER);
            float[] x = new float[arr.length];
            for (int i = 0; i < arr.length; i++) x[i] = (float) arr[i];

            // 2) 构建输入张量 [1, N]
            OnnxTensor t = OnnxTensor.createTensor(env, FloatBuffer.wrap(x), new long[]{1, x.length});
            try {
                Map<String, OnnxTensorLike> input = new HashMap<String, OnnxTensorLike>(1);
                String inName = session.getInputNames().iterator().next();
                input.put(inName, t);

                // 3) 推理
                OrtSession.Result r = null;
                try {
                    r = session.run(input);

                    // 4) 取输出：显式名优先；没有就第一个输出
                    OnnxValue v = null;
                    if (this.outputName != null) {
                        Optional<OnnxValue> opt = r.get(this.outputName);
                        if (opt != null && opt.isPresent()) v = opt.get();
                    }
                    if (v == null) v = r.get(0);

                    // 5) 解析张量 -> 分数
                    double raw = 0.0;
                    if (v instanceof OnnxTensor) {
                        OnnxTensor ot = (OnnxTensor) v;
                        Object val = ot.getValue();
                        if (val instanceof float[][])      raw = ((float[][]) val)[0][0];
                        else if (val instanceof float[])   raw = ((float[]) val)[0];
                        else if (val instanceof double[][])raw = ((double[][]) val)[0][0];
                        else if (val instanceof double[])  raw = ((double[]) val)[0];
                        else raw = 0.0;
                    }

                    // 6) 归一化与方向（OCSVM decision_function >0 正常、<0 异常 → 需翻转）
                    double s = raw;
                    if (flipDecision) {
                        // 将 decision_function 映射到风险分：越小越危险 → 风险高
                        // 这里简单用 sigmoid(-x)；你也可以做分位/Platt 标定
                        s = 1.0 / (1.0 + Math.exp(raw)); // sigmoid(-raw)
                    }
                    return clamp(s);
                } finally {
                    if (r != null) r.close();
                }
            } finally {
                t.close();
            }
        } catch (Exception e) {
            log.error("ONNX模型推理失败，特征向量: {}", fv.values, e);
            //todo 发送告警（钉钉/企微/邮件）

            // 推理异常保守返回低分（避免误杀），可在日志侧告警
            return 0.0;
        }
    }

    private static double clamp(double v){ return Math.max(0.0, Math.min(1.0, v)); }

    @Override
    public void close() throws Exception {
        if (session != null) session.close();
        if (env != null) env.close();
    }
}
